import torch
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
"""
Consider your input u_n and u_next as [BatchSize,1,Height,Width]

"""



def loss_be(u_n, u_next, deltaT, domain,epsilon,type='BE'):
    # case2
    if type=='FE':
         residual = u_next-u_n-deltaT*(epsilon*epsilon*FiniteDiff_LB(u_n, domain) + u_n-u_n**3)
    elif type=='BE':
        residual = (u_next-u_n)-deltaT*(epsilon*epsilon*FiniteDiff_LB(u_next, domain) + u_next-u_next**3)
    elif type=='CN':
        residual = (u_next - u_n)-deltaT*(epsilon * epsilon * (FiniteDiff_LB(u_n, domain)+FiniteDiff_LB(u_next, domain))/2 + (u_next+u_n)/2 - (u_next ** 3+u_n**3)/2)


    return residual




def FiniteDiff_LB(x, domain):
    # consider a left bottom 4X4 matrix as input given a 5X5 image. Periodic boundary condition is applied
    # x = x[:,:,1:,0:-1]
    #Note this works for LeftBottom Part only!
    h_x = x.size()[2]
    w_x = x.size()[3]
    lx = domain[1] - domain[0]
    ly = domain[3] - domain[2]
    dx = lx / (w_x)
    dy = ly / (h_x)

    # FD = torch.zeros_like(x)
    FD_right = torch.zeros_like(x)
    FD_left = torch.zeros_like(x)
    FD_top = torch.zeros_like(x)
    FD_bot = torch.zeros_like(x)
    FD_right[:,:,:,w_x-1] = x[:,:,:,0]
    FD_right[:,:,:,0:w_x-1] = x[:,:,:,1:w_x]
    FD_left[:,:,:,0] = x[:,:,:,w_x-1]
    FD_left[:,:,:,1:w_x] = x[:,:,:,0:w_x-1]
    FD_top[:,:,0,:] = x[:,:,h_x-1,:]
    FD_top[:,:,1:h_x,:] = x[:,:,0:h_x-1,:]
    FD_bot[:,:,h_x-1,:] = x[:,:,0,:]
    FD_bot[:,:,0:h_x-1,:] = x[:,:,1:h_x,:]

    FD = (FD_right + FD_left - 2*x)/(dx**2) + (FD_top + FD_bot - 2*x)/(dy**2)
    return FD

def BioFiniteDiff_LB(x, domain):
    # consider a left bottom 4X4 matrix as input given a 5X5 image. Periodic boundary condition is applied
    # x = x[:,:,1:,0:-1]
    #Note this works for LeftBottom Part only!
    h_x = x.size()[2]
    w_x = x.size()[3]
    lx = domain[1] - domain[0]
    ly = domain[3] - domain[2]
    dx = lx / (w_x)
    dy = ly / (h_x)

    # FD = torch.zeros_like(x)
    FD_right1 = torch.zeros_like(x)
    FD_right2 = torch.zeros_like(x)
    FD_left1 = torch.zeros_like(x)
    FD_left2 = torch.zeros_like(x)
    FD_top1 = torch.zeros_like(x)
    FD_top2 = torch.zeros_like(x)
    FD_bot1 = torch.zeros_like(x)
    FD_bot2 = torch.zeros_like(x)

    FD_diag1 = torch.zeros_like(x)
    FD_diag2 = torch.zeros_like(x)
    FD_diag3 = torch.zeros_like(x)
    FD_diag4 = torch.zeros_like(x)

    FD_right1[:,:,:,w_x-1] = -4*x[:,:,:,0]
    FD_right1[:,:,:,0:w_x-1] = -4*x[:,:,:,1:w_x]
    FD_right2[:, :, :, w_x - 1] = x[:, :, :, 1]
    FD_right2[:, :, :, w_x - 2] = x[:, :, :, 0]
    FD_right2[:, :, :, 0:w_x - 2] = x[:, :, :, 2:w_x]

    FD_left1[:,:,:,0] = -4*x[:,:,:,w_x-1]
    FD_left1[:,:,:,1:w_x] = -4*x[:,:,:,0:w_x-1]
    FD_left2[:, :, :, 0] = x[:, :, :, w_x - 2]
    FD_left2[:, :, :, 1] = x[:, :, :, w_x - 1]
    FD_left2[:, :, :, 2:w_x] = x[:, :, :, 0:w_x - 2]

    FD_top1[:,:,0,:] = -4 * x[:,:,h_x-1,:]
    FD_top1[:,:,1:h_x,:] = -4 * x[:,:,0:h_x-1,:]
    FD_top2[:, :, 0, :] = x[:, :, h_x - 2, :]
    FD_top2[:, :, 1, :] = x[:, :, h_x - 1, :]
    FD_top2[:, :, 2:h_x, :] = x[:, :, 0:h_x - 2, :]

    FD_bot1[:,:,h_x-1,:] = -4 * x[:,:,0,:]
    FD_bot1[:,:,0:h_x-1,:] = -4 * x[:,:,1:h_x,:]
    FD_bot2[:, :, h_x - 1, :] = x[:, :, 1, :]
    FD_bot2[:, :, h_x - 2, :] = x[:, :, 0, :]
    FD_bot2[:, :, 0:h_x - 2, :] = x[:, :, 2:h_x, :]

    FD_diag1[:,:,0:h_x-1,0:w_x-1] = 2 * x[:,:,1:h_x,1:w_x]
    FD_diag1[:,:,h_x-1,0:w_x-1] = 2 * x[:,:,0,1:w_x]
    FD_diag1[:, :, 0:h_x-1, w_x-1] = 2 * x[:, :, 1:h_x, 0]
    FD_diag1[:, :, h_x-1, w_x-1] = 2 * x[:, :, 0, 0]

    FD_diag2[:, :, 1:h_x, 1:w_x] = 2 * x[:, :, 0:h_x-1, 0:w_x-1]
    FD_diag2[:, :, 0, 1:w_x] = 2 * x[:, :, h_x-1, 0:w_x-1]
    FD_diag2[:, :, 1:h_x, 0] = 2 * x[:, :, 0:h_x-1, w_x-1]
    FD_diag2[:, :, 0, 0] = 2 * x[:, :, h_x-1, w_x-1]

    FD_diag3[:, :, 1:h_x, 0:w_x-1] = 2 * x[:, :, 0:h_x -1, 1:w_x]
    FD_diag3[:, :, 0, 0:w_x-1] = 2 * x[:, :, h_x-1, 1:w_x]
    FD_diag3[:, :, 1:h_x, w_x-1] = 2 * x[:, :, 0:h_x - 1, 0]
    FD_diag3[:, :, 0, w_x-1] = 2 * x[:, :, h_x-1, 0]

    FD_diag4[:, :, 0:h_x-1, 1:w_x] = 2 * x[:, :, 1:h_x, 0:w_x-1]
    FD_diag4[:, :, h_x-1, 1:w_x] = 2 * x[:, :, 0, 0:w_x-1]
    FD_diag4[:, :, 0:h_x - 1, 0] = 2 * x[:, :, 1:h_x, w_x-1]
    FD_diag4[:, :, h_x-1, 0] = 2 * x[:, :, 0, w_x-1]

    FD = (FD_right1 + FD_right2 +FD_left1 +FD_left2 + 6*x)/(dx**4) + (FD_top1 + FD_top2 +FD_bot1 +FD_bot2 + 6*x)/(dy**4)+(FD_diag1+FD_diag2+FD_diag3+FD_diag4+FD_right1+FD_left1+FD_top1+FD_bot1+8*x)/((dx**2)*(dy**2))
    return FD


def integral(x):
    h_x = x.size()[2]
    w_x = x.size()[3]
    a = x.mean([2, 3]).unsqueeze(2).unsqueeze(3)
    b = a.expand(-1, -1, h_x, w_x)
    return b

def mass(x):
    h_x = x.size()[2]
    w_x = x.size()[3]
    a = x.mean([2, 3]).unsqueeze(2).unsqueeze(3)
    b = a.expand(-1, -1, h_x, w_x)
    return a
def mass_engy(U,epsilon,domain):
    mass = 0
    egy1 = 0
    egy2 = 0
    N = U.size()[2]
    M = U.size()[3]
    lx = domain[1] - domain[0]
    ly = domain[3] - domain[2]
    hx = lx / (N-1)
    hy = ly / (M-1)

    for i in range(N - 1):
        for j in range(M - 1):
            uu = 0.25 * (U[:,:,i, j] + U[:,:,i + 1, j] + U[:,:,i, j + 1] + U[:,:,i + 1, j + 1])
            mass = mass + uu
            egy1 = egy1 + (uu * uu - 1) ** 2
            udx = 0.5 * (U[:,:,i + 1, j] - U[:,:,i, j] + U[:,:,i + 1, j + 1] - U[:,:,i, j + 1]) / hx
            udy = 0.5 * (U[:,:,i, j + 1] - U[:,:,i, j] + U[:,:,i + 1, j + 1] - U[:,:,i + 1, j]) / hy
            egy2 = egy2 + (udx * udx + udy * udy)

    mass = mass * hx * hy
    egy1 = 0.25 * egy1
    egy2 = 0.5 * egy2 * epsilon ** 2
    egy = (egy1 + egy2) * hx * hy
    return mass, egy
# def forward(x, domain):
#     batch_size=x.size()[0]
#     h_x=x.size()[2]
#     w_x=x.size()[3]
#     lx =domain[1]-domain[0]
#     ly = domain[3] - domain[2]
#
#     deltax = lx/(w_x-1)
#     deltay = ly/(h_x-1)
#
#     h_tv =  (x[:,:,2:h_x,:]+x[:,:,:h_x-2,:]-2*(x[:,:,1:h_x-1,:]))/(deltay*deltay)
#     w_tv =  (x[:,:,:,2:w_x]+x[:,:,:,:w_x-2]-2*(x[:,:,:,1:w_x-1]))/(deltax*deltax)
#     tv= h_tv[:,:,:,1:w_x-1]+w_tv[:,:,1:h_x-1,:]
#     return tv
#
# def leftboundary(x,domain):
#     batch_size = x.size()[0]
#     h_x = x.size()[2]
#     w_x = x.size()[3]
#     lx = domain[1] - domain[0]
#     ly = domain[3] - domain[2]
#
#     deltax = lx / (w_x - 1)
#     deltay = ly / (h_x - 1)
#
#     le_h_btv=(x[:,:,2:h_x,[0]]+x[:,:,:h_x-2,[0]]-2*x[:,:,1:h_x-1,[0]])/(deltay*deltay)
#     le_w_btv=(x[:,:,1:h_x-1,[1]]+x[:,:,1:h_x-1,[w_x-2]]-2*(x[:,:,1:h_x-1,[0]]))/(deltax*deltax)
#     return le_h_btv+le_w_btv
#
# def lowerboundary(x,domain):
#     batch_size = x.size()[0]
#     h_x = x.size()[2]
#     w_x = x.size()[3]
#     lx = domain[1] - domain[0]
#     ly = domain[3] - domain[2]
#
#     deltax = lx / (w_x - 1)
#     deltay = ly / (h_x - 1)
#     lo_h_btv=(x[:,:,[h_x-2],1:w_x-1]+x[:,:,[1],1:w_x-1]-2*x[:,:,[h_x-1],1:w_x-1])/(deltay*deltay)
#     lo_w_btv=(x[:,:,[h_x-1],2:w_x]+x[:,:,[h_x-1],:w_x-2]-2*x[:,:,[h_x-1],1:w_x-1])/(deltax*deltax)
#     return lo_h_btv+lo_w_btv
#
# def corner(x, domain):
#     batch_size = x.size()[0]
#     h_x = x.size()[2]
#     w_x = x.size()[3]
#     lx = domain[1] - domain[0]
#     ly = domain[3] - domain[2]
#
#     deltax = lx / (w_x - 1)
#     deltay = ly / (h_x - 1)
#     c_h_tv=(x[:,:,h_x-2,0]+x[:,:,1,0]-2*x[:,:,h_x-1,0])/(deltay*deltay)
#     c_w_tv=(x[:,:,h_x-1,1]+x[:,:,h_x-1,w_x-2]-2*x[:,:,h_x-1,0])/(deltax*deltax)
#     c_tv=c_h_tv+c_w_tv
#     return  c_tv
#
# def FiniteDiff(x, domain):
#     batch_size = x.size()[0]
#     h_x = x.size()[2]
#     w_x = x.size()[3]
#     FD = torch.zeros_like(x)
#     FD[:,:,1:h_x-1,1:w_x-1]=forward(x, domain)
#     FD[:,:,1:h_x-1,[0]]=leftboundary(x,domain)
#     FD[:, :, [h_x - 1], 1:w_x - 1] = lowerboundary(x,domain)
#     FD[:,:,h_x-1,0]=corner(x,domain)
#     FD[:,:,0,0]=corner(x,domain)
#     FD[:, :, h_x-1, w_x-1] = corner(x, domain)
#     FD[:,:,:h_x-1,[w_x-1]]=FD[:,:,:h_x-1,[0]]
#     FD[:,:,[0],:w_x-1]=FD[:,:,[h_x-1],:w_x-1]
#     return FD
#



# h_x=x.size()[2]
# w_x=x.size()[3]
#
# FD_right1 = torch.zeros_like(x)
# FD_right2 = torch.zeros_like(x)
# FD_left1 = torch.zeros_like(x)
# FD_left2 = torch.zeros_like(x)
# FD_top1 = torch.zeros_like(x)
# FD_top2 = torch.zeros_like(x)
# FD_bot1 = torch.zeros_like(x)
# FD_bot2 = torch.zeros_like(x)
#
# FD_diag1 = torch.zeros_like(x)
# FD_diag2 = torch.zeros_like(x)
# FD_diag3 = torch.zeros_like(x)
# FD_diag4 = torch.zeros_like(x)
#
# FD_right1[:,:,:,w_x-1] = -4*x[:,:,:,0]
# FD_right1[:,:,:,0:w_x-1] = -4*x[:,:,:,1:w_x]
# FD_right2[:, :, :, w_x - 1] = x[:, :, :, 1]
# FD_right2[:, :, :, w_x - 2] = x[:, :, :, 0]
# FD_right2[:, :, :, 0:w_x - 2] = x[:, :, :, 2:w_x]
#
# FD_left1[:,:,:,0] = -4*x[:,:,:,w_x-1]
# FD_left1[:,:,:,1:w_x] = -4*x[:,:,:,0:w_x-1]
# FD_left2[:, :, :, 0] = x[:, :, :, w_x - 2]
# FD_left2[:, :, :, 1] = x[:, :, :, w_x - 1]
# FD_left2[:, :, :, 2:w_x] = x[:, :, :, 0:w_x - 2]
#
# FD_top1[:,:,0,:] = -4 * x[:,:,h_x-1,:]
# FD_top1[:,:,1:h_x,:] = -4 * x[:,:,0:h_x-1,:]
# FD_top2[:, :, 0, :] = x[:, :, h_x - 2, :]
# FD_top2[:, :, 1, :] = x[:, :, h_x - 1, :]
# FD_top2[:, :, 2:h_x, :] = x[:, :, 0:h_x - 2, :]
#
# FD_bot1[:,:,h_x-1,:] = -4 * x[:,:,0,:]
# FD_bot1[:,:,0:h_x-1,:] = -4 * x[:,:,1:h_x,:]
# FD_bot2[:, :, h_x - 1, :] = x[:, :, 1, :]
# FD_bot2[:, :, h_x - 2, :] = x[:, :, 0, :]
# FD_bot2[:, :, 0:h_x - 2, :] = x[:, :, 2:h_x, :]
#
# FD_diag1[:,:,0:h_x-1,0:w_x-1] = 2 * x[:,:,1:h_x,1:w_x]
# FD_diag1[:,:,h_x-1,0:w_x-1] = 2 * x[:,:,0,1:w_x]
# FD_diag1[:, :, 0:h_x-1, w_x-1] = 2 * x[:, :, 1:h_x, 0]
# FD_diag1[:, :, h_x-1, w_x-1] = 2 * x[:, :, 0, 0]
#
# FD_diag2[:, :, 1:h_x, 1:w_x] = 2 * x[:, :, 0:h_x-1, 0:w_x-1]
# FD_diag2[:, :, 0, 1:w_x] = 2 * x[:, :, h_x-1, 0:w_x-1]
# FD_diag2[:, :, 1:h_x, 0] = 2 * x[:, :, 0:h_x-1, w_x-1]
# FD_diag2[:, :, 0, 0] = 2 * x[:, :, h_x-1, w_x-1]
#
# FD_diag3[:, :, 1:h_x, 0:w_x-1] = 2 * x[:, :, 0:h_x -1, 1:w_x]
# FD_diag3[:, :, 0, 0:w_x-1] = 2 * x[:, :, h_x-1, 1:w_x]
# FD_diag3[:, :, 1:h_x, w_x-1] = 2 * x[:, :, 0:h_x - 1, 0]
# FD_diag3[:, :, 0, w_x-1] = 2 * x[:, :, h_x-1, 0]
#
# FD_diag4[:, :, 0:h_x-1, 1:w_x] = 2 * x[:, :, 1:h_x, 0:w_x-1]
# FD_diag4[:, :, h_x-1, 1:w_x] = 2 * x[:, :, 0, 0:w_x-1]
# FD_diag4[:, :, 0:h_x - 1, 0] = 2 * x[:, :, 1:h_x, w_x-1]
# FD_diag4[:, :, h_x-1, 0] = 2 * x[:, :, 0, w_x-1]
# print(2*x)
# print(FD_diag4)
